package org.bing.zion.filter.header;

import java.io.IOException;
import java.nio.ByteBuffer;

import org.bing.zion.core.ProtocolFilterAdapter;
import org.bing.zion.core.ProtocolFilterChain;
import org.bing.zion.core.Session;


public class HandShakeFilter extends ProtocolFilterAdapter {
	private static final String SESSION_KEY_HEADER = "KEY.FILTER.SHAKE";
	private static final String SESSION_KEY_CURRSIZE = "KEY.FILTER.SHAKE.CURRSIZE";
	private static final String SESSION_KEY_CURRFRAG = "KEY.FILTER.SHAKE.CURRFRAG";
	private static final int HEADER_MAX_SIZE = 1024;
	private static final byte HEADER_END_BYTE = '\n';

	/**
	 * HEADER: WATCHER/1.1 b77b09a8ec8aac53259b4fe312bc1cc3 [/n]
	 */

	@Override
	public void messageReceived(ProtocolFilterChain chain, Session session, Object msg) throws IOException {
		String header = (String) session.getAttribute(SESSION_KEY_HEADER);
		if (header != null || !(msg instanceof ByteBuffer)) {
			super.messageReceived(chain, session, msg);
		} else if (msg instanceof ByteBuffer) {
			process(chain, session, msg);
		}
	}

	private void process(ProtocolFilterChain chain, Session session, Object msg) throws IOException {
		ByteBuffer buf = (ByteBuffer) msg;
		Integer size = (Integer) session.getAttribute(SESSION_KEY_CURRSIZE);
		if (size == null) {
			if (!doRead(session, buf, new byte[0])) {
				super.messageReceived(chain, session, buf);
			}
		} else if (size < HEADER_MAX_SIZE) {
			byte[] frag = (byte[]) session.getAttribute(SESSION_KEY_CURRFRAG);
			if (!doRead(session, buf, frag)) {
				super.messageReceived(chain, session, msg);
			}
		} else { // not header ( old version )
			super.messageReceived(chain, session, msg);
		}
	}

	private boolean doRead(Session session, ByteBuffer buf, byte[] frag) {
		int len = Math.min(HEADER_MAX_SIZE - frag.length, buf.position());
		boolean flag = false;
		byte[] hbuf = null;
		for (int i = 0; i < len; i++) {
			byte b = buf.get(i);
			if (b == HEADER_END_BYTE) {
				flag = true;
				hbuf = new byte[i + 1];
				buf.flip();
				buf.get(hbuf);
				break;
			}
		}
		if (flag) { // has header
			String header = new String(hbuf).trim();
			logger.info("Header:" + header);
			session.setAttribute(SESSION_KEY_HEADER, header);
			buf = buf.compact();
			return false;
		} else {
			if (len + frag.length == HEADER_MAX_SIZE) { // not header, old version
				String header = "WATCHER/1.0 b77b09a8ec8aac53259b4fe312bc1cc3";
				logger.info("HD:" + header);
				session.setAttribute(SESSION_KEY_HEADER, header);
				byte[] tmp = new byte[HEADER_MAX_SIZE];
				System.arraycopy(frag, 0, tmp, 0, frag.length);
				System.arraycopy(buf.array(), 0, tmp, frag.length, buf.limit()); 
				buf = ByteBuffer.wrap(tmp);
				return false;
			} else { // update current header size
				int currSize = frag.length + len;
				session.setAttribute(SESSION_KEY_CURRSIZE, currSize);
				byte[] tmp = new byte[currSize];
				System.arraycopy(frag, 0, tmp, 0, frag.length);
				System.arraycopy(hbuf, 0, tmp, frag.length, hbuf.length);
				session.setAttribute(SESSION_KEY_CURRFRAG, tmp);
				return true;
			}
		}
	}
}
